# -*- coding: utf-8 -*-
"""
Created on Thu Aug  4 18:43:38 2022

@author: Shang'xiang

Stay Hungry Stay Foolish
"""

'''
此程序用于走时CT的GN反演
'''

# 初始化，读取数据
import ttcrpy.rgrid as rg
import numpy as np
from scipy.sparse.linalg import cg

import matplotlib.pyplot as plt

data = np.loadtxt('G02_FSM.txt')

srcs = data[:,0:2]
rcvs = data[:,2:4]
tobs = data[:,4]

# 给定初始模型
# 创建网格
x = np.arange(0,31.0)
y = np.arange(0,21.0)

# 创建速度模型
v = 2000*np.ones((x.size-1,y.size-1))


# 离散网格
grid = rg.Grid2d(x, y, cell_slowness=True, method = 'SPM')

# 速度转换为慢度
slowness = 1./v

# =============================================================================
# # tcal,LL = grid.raytrace(srcs, rcvs, slowness, compute_L = True)
# 
# # A = LL.todense()
# 
# # I = np.eye(600)
# 
# # zuo = A.T*A + 0.618*I
# # deltat = tobs - tcal
# # you = A.T.dot(deltat)
# # you = you.T
# 
# # deltam,info = cg(zuo,you)
# =============================================================================

I = np.eye(600)

# =============================================================================
# # CT正演
# =============================================================================
def ctforward(slowness):
    tcal,LL = grid.raytrace(srcs, rcvs, slowness, compute_L = True)
    A = LL.todense()
    zuo = A.T*A + 0.618*I
    deltat = tobs - tcal
    you = A.T.dot(deltat)
    you = you.T
    deltam,info = cg(zuo,you)
    return deltam

# =============================================================================
# # 第一次计算
# # 计算模型修改量
deltam = ctforward(slowness)
deltam = deltam.reshape(30,20)
# 
# # 修改慢度
slowness = slowness + deltam
# =============================================================================

# 开始循环
for i in range(100):
    deltam = ctforward(slowness)
    deltam = deltam.reshape(30,20)
    slowness = slowness + deltam
    v = 1./slowness
    # 画模型图
    fig, ax = plt.subplots()
    cs = plt.pcolor(v,cmap='jet',edgecolor='w')
    plt.colorbar()
    plt.show()












    




